{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "d9a8ac2c",
   "metadata": {},
   "source": [
    "## Preparation\n",
    "\n",
    "### Install packages\n",
    "\n",
    "Make sure you have installed the newest python packages:\n",
    "\n",
    "| package |\n",
    "| -- |\n",
    "| towhee |\n",
    "| transformers |\n",
    "| datasets |\n",
    "| evaluate |\n",
    "| scikit-learn |\n",
    "| torch |\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "2d6b1c37",
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "# ! python -m pip install torch towhee transformers datasets evaluate scikit-learn"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "938065a8",
   "metadata": {},
   "source": [
    "You can try the following block code to make sure your environment is valid."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "13d99adf",
   "metadata": {},
   "outputs": [],
   "source": [
    "# try this line\n",
    "from torch.distributed import ProcessGroup"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "406acc3a",
   "metadata": {},
   "source": [
    "If you get a Error like list: cannot import name 'ProcessGroup' from 'torch.distributed', please refer [this issue](https://github.com/pytorch/pytorch/issues/68385#issuecomment-1332607943) and install the newest pytorch version."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "014508ce",
   "metadata": {},
   "source": [
    "## Fine-tune BERT on Masked Language Modeling task"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3ce420fa",
   "metadata": {},
   "source": [
    "[Bidirectional Encoder Representations from Transformers (BERT)](https://arxiv.org/abs/1810.04805)  is a transformer-based machine learning technique for natural language processing (NLP) pre-training developed by Google.BERT was pretrained on two tasks: language modeling and next sentence prediction. As a result of the training process, BERT learns contextual embeddings for words. After pretraining, which is computationally expensive, BERT can be finetuned with fewer resources on smaller datasets to optimize its performance on specific tasks."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0af11b72",
   "metadata": {},
   "source": [
    "### Instantiate operator\n",
    "We can instantiate a towhee operator containing the [`bert-base-uncased`](https://huggingface.co/bert-base-uncased) model.  \n",
    "\n",
    "**Note**: By default, we initialize `text_embedding.transformers` to get embedding for every token when inference. If you want to get a sentence embedding rather than embedding for every token, please use `sentence_embedding.transformers` operator."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "5c46a4e6",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Some weights of the model checkpoint at bert-base-uncased were not used when initializing BertModel: ['cls.predictions.transform.LayerNorm.bias', 'cls.predictions.transform.LayerNorm.weight', 'cls.predictions.decoder.weight', 'cls.seq_relationship.weight', 'cls.predictions.transform.dense.bias', 'cls.predictions.transform.dense.weight', 'cls.seq_relationship.bias', 'cls.predictions.bias']\n",
      "- This IS expected if you are initializing BertModel from the checkpoint of a model trained on another task or with another architecture (e.g. initializing a BertForSequenceClassification model from a BertForPreTraining model).\n",
      "- This IS NOT expected if you are initializing BertModel from the checkpoint of a model that you expect to be exactly identical (initializing a BertForSequenceClassification model from a BertForSequenceClassification model).\n"
     ]
    }
   ],
   "source": [
    "import towhee\n",
    "\n",
    "bert_op = towhee.ops.text_embedding.transformers(model_name='bert-base-uncased').get_op()\n",
    "# or (if you want to get a sentence embedding rather than embedding for every token, please use sentence_embedding operator)\n",
    "# bert_op = towhee.ops.sentence_embedding.transformers(model_name='bert-base-uncased').get_op()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e26a6057",
   "metadata": {},
   "source": [
    "This operator can embed a sentence using the pretrained models's output of the last layer before the model's head. The embedding shape is `([input token num], [model dim])`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "fd11b04b",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(array([[-0.29221347,  0.16306752,  0.20236042, ..., -0.21968302,\n",
       "         -0.20541406,  0.7011004 ],\n",
       "        [-0.61254054,  0.55492014,  1.4136117 , ...,  0.20618206,\n",
       "          0.3001338 ,  1.0520668 ],\n",
       "        [ 0.62186277,  0.40705413,  1.0701158 , ..., -0.04220748,\n",
       "          0.25006238,  0.68540764],\n",
       "        ...,\n",
       "        [-0.6860294 ,  0.04325684,  1.0766182 , ...,  0.25660843,\n",
       "         -0.36851028,  0.43962553],\n",
       "        [-0.5295304 , -0.33464265,  0.10741595, ...,  0.86484456,\n",
       "          0.02151933, -0.08432791],\n",
       "        [ 0.79523426,  0.26699936,  0.03003302, ...,  0.24721493,\n",
       "         -0.55750215, -0.16443911]], dtype=float32),\n",
       " (9, 768))"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "embedding = bert_op('Hello world, hello every one.')\n",
    "embedding, embedding.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "15a37c75",
   "metadata": {},
   "source": [
    "### Start training\n",
    "We just specify two args dict, and run method `train()` with `task='mlm'` to start training. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "bfd6bd92",
   "metadata": {},
   "outputs": [],
   "source": [
    "data_args = {\n",
    "    'dataset_name': 'wikitext',\n",
    "    'dataset_config_name': 'wikitext-2-raw-v1',\n",
    "}\n",
    "training_args = {\n",
    "    'num_train_epochs': 3, # you can add epoch number to get a better metric.\n",
    "    'per_device_train_batch_size': 8,\n",
    "    'per_device_eval_batch_size': 8,\n",
    "    'do_train': True,\n",
    "    'do_eval': True,\n",
    "    'output_dir': './tmp/test-mlm',\n",
    "    'overwrite_output_dir': True\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "6fe2a7b7",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Some weights of the model checkpoint at bert-base-uncased were not used when initializing BertForMaskedLM: ['cls.seq_relationship.bias', 'cls.seq_relationship.weight']\n",
      "- This IS expected if you are initializing BertForMaskedLM from the checkpoint of a model trained on another task or with another architecture (e.g. initializing a BertForSequenceClassification model from a BertForPreTraining model).\n",
      "- This IS NOT expected if you are initializing BertForMaskedLM from the checkpoint of a model that you expect to be exactly identical (initializing a BertForSequenceClassification model from a BertForSequenceClassification model).\n",
      "2022-12-20 09:05:26,097 - 140132563211072 - train_mlm_with_hf_trainer.py-train_mlm_with_hf_trainer:159 - WARNING: Process rank: -1, device: cuda:0, n_gpu: 8, distributed training: False, 16-bits training: False\n",
      "2022-12-20 09:05:26,098 - 140132563211072 - train_mlm_with_hf_trainer.py-train_mlm_with_hf_trainer:164 - INFO: Training/evaluation parameters TrainingArguments(\n",
      "_n_gpu=8,\n",
      "adafactor=False,\n",
      "adam_beta1=0.9,\n",
      "adam_beta2=0.999,\n",
      "adam_epsilon=1e-08,\n",
      "auto_find_batch_size=False,\n",
      "bf16=False,\n",
      "bf16_full_eval=False,\n",
      "data_seed=None,\n",
      "dataloader_drop_last=False,\n",
      "dataloader_num_workers=0,\n",
      "dataloader_pin_memory=True,\n",
      "ddp_bucket_cap_mb=None,\n",
      "ddp_find_unused_parameters=None,\n",
      "ddp_timeout=1800,\n",
      "debug=[],\n",
      "deepspeed=None,\n",
      "disable_tqdm=False,\n",
      "do_eval=True,\n",
      "do_predict=False,\n",
      "do_train=True,\n",
      "eval_accumulation_steps=None,\n",
      "eval_delay=0,\n",
      "eval_steps=None,\n",
      "evaluation_strategy=no,\n",
      "fp16=False,\n",
      "fp16_backend=auto,\n",
      "fp16_full_eval=False,\n",
      "fp16_opt_level=O1,\n",
      "fsdp=[],\n",
      "fsdp_min_num_params=0,\n",
      "fsdp_transformer_layer_cls_to_wrap=None,\n",
      "full_determinism=False,\n",
      "gradient_accumulation_steps=1,\n",
      "gradient_checkpointing=False,\n",
      "greater_is_better=None,\n",
      "group_by_length=False,\n",
      "half_precision_backend=auto,\n",
      "hub_model_id=None,\n",
      "hub_private_repo=False,\n",
      "hub_strategy=every_save,\n",
      "hub_token=<HUB_TOKEN>,\n",
      "ignore_data_skip=False,\n",
      "include_inputs_for_metrics=False,\n",
      "jit_mode_eval=False,\n",
      "label_names=None,\n",
      "label_smoothing_factor=0.0,\n",
      "learning_rate=5e-05,\n",
      "length_column_name=length,\n",
      "load_best_model_at_end=False,\n",
      "local_rank=-1,\n",
      "log_level=passive,\n",
      "log_level_replica=passive,\n",
      "log_on_each_node=True,\n",
      "logging_dir=./tmp/test-mlm/runs/Dec20_09-05-26_zilliz,\n",
      "logging_first_step=False,\n",
      "logging_nan_inf_filter=True,\n",
      "logging_steps=500,\n",
      "logging_strategy=steps,\n",
      "lr_scheduler_type=linear,\n",
      "max_grad_norm=1.0,\n",
      "max_steps=-1,\n",
      "metric_for_best_model=None,\n",
      "mp_parameters=,\n",
      "no_cuda=False,\n",
      "num_train_epochs=3,\n",
      "optim=adamw_hf,\n",
      "optim_args=None,\n",
      "output_dir=./tmp/test-mlm,\n",
      "overwrite_output_dir=True,\n",
      "past_index=-1,\n",
      "per_device_eval_batch_size=8,\n",
      "per_device_train_batch_size=8,\n",
      "prediction_loss_only=False,\n",
      "push_to_hub=False,\n",
      "push_to_hub_model_id=None,\n",
      "push_to_hub_organization=None,\n",
      "push_to_hub_token=<PUSH_TO_HUB_TOKEN>,\n",
      "ray_scope=last,\n",
      "remove_unused_columns=True,\n",
      "report_to=[],\n",
      "resume_from_checkpoint=None,\n",
      "run_name=./tmp/test-mlm,\n",
      "save_on_each_node=False,\n",
      "save_steps=500,\n",
      "save_strategy=steps,\n",
      "save_total_limit=None,\n",
      "seed=42,\n",
      "sharded_ddp=[],\n",
      "skip_memory_metrics=True,\n",
      "tf32=None,\n",
      "torchdynamo=None,\n",
      "tpu_metrics_debug=False,\n",
      "tpu_num_cores=None,\n",
      "use_ipex=False,\n",
      "use_legacy_prediction_loop=False,\n",
      "use_mps_device=False,\n",
      "warmup_ratio=0.0,\n",
      "warmup_steps=0,\n",
      "weight_decay=0.0,\n",
      "xpu_backend=None,\n",
      ")\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "train mlm with hugging face transformers trainer\n",
      "**** DataTrainingArguments ****\n",
      "- dataset_name \n",
      "  default: None \n",
      "  metadata_dict: {'help': 'The name of the dataset to use (via the datasets library).'} \n",
      "\n",
      "- dataset_config_name \n",
      "  default: None \n",
      "  metadata_dict: {'help': 'The configuration name of the dataset to use (via the datasets library).'} \n",
      "\n",
      "- train_file \n",
      "  default: None \n",
      "  metadata_dict: {'help': 'The input training data file (a text file).'} \n",
      "\n",
      "- validation_file \n",
      "  default: None \n",
      "  metadata_dict: {'help': 'An optional input evaluation data file to evaluate the perplexity on (a text file).'} \n",
      "\n",
      "- overwrite_cache \n",
      "  default: False \n",
      "  metadata_dict: {'help': 'Overwrite the cached training and evaluation sets'} \n",
      "\n",
      "- validation_split_percentage \n",
      "  default: 5 \n",
      "  metadata_dict: {'help': \"The percentage of the train set used as validation set in case there's no validation split\"} \n",
      "\n",
      "- max_seq_length \n",
      "  default: None \n",
      "  metadata_dict: {'help': 'The maximum total input sequence length after tokenization. Sequences longer than this will be truncated.'} \n",
      "\n",
      "- preprocessing_num_workers \n",
      "  default: None \n",
      "  metadata_dict: {'help': 'The number of processes to use for the preprocessing.'} \n",
      "\n",
      "- mlm_probability \n",
      "  default: 0.15 \n",
      "  metadata_dict: {'help': 'Ratio of tokens to mask for masked language modeling loss'} \n",
      "\n",
      "- line_by_line \n",
      "  default: False \n",
      "  metadata_dict: {'help': 'Whether distinct lines of text in the dataset are to be handled as distinct sequences.'} \n",
      "\n",
      "- pad_to_max_length \n",
      "  default: False \n",
      "  metadata_dict: {'help': 'Whether to pad all samples to `max_seq_length`. If False, will pad the samples dynamically when batching to the maximum length in the batch.'} \n",
      "\n",
      "- max_train_samples \n",
      "  default: None \n",
      "  metadata_dict: {'help': 'For debugging purposes or quicker training, truncate the number of training examples to this value if set.'} \n",
      "\n",
      "- max_eval_samples \n",
      "  default: None \n",
      "  metadata_dict: {'help': 'For debugging purposes or quicker training, truncate the number of evaluation examples to this value if set.'} \n",
      "\n",
      "**** TrainingArguments ****\n",
      "- output_dir \n",
      "  default: <dataclasses._MISSING_TYPE object at 0x7f73254f2760> \n",
      "  metadata_dict: {'help': 'The output directory where the model predictions and checkpoints will be written.'} \n",
      "\n",
      "- overwrite_output_dir \n",
      "  default: False \n",
      "  metadata_dict: {'help': 'Overwrite the content of the output directory. Use this to continue training if output_dir points to a checkpoint directory.'} \n",
      "\n",
      "- do_train \n",
      "  default: False \n",
      "  metadata_dict: {'help': 'Whether to run training.'} \n",
      "\n",
      "- do_eval \n",
      "  default: False \n",
      "  metadata_dict: {'help': 'Whether to run eval on the dev set.'} \n",
      "\n",
      "- do_predict \n",
      "  default: False \n",
      "  metadata_dict: {'help': 'Whether to run predictions on the test set.'} \n",
      "\n",
      "- evaluation_strategy \n",
      "  default: no \n",
      "  metadata_dict: {'help': 'The evaluation strategy to use.'} \n",
      "\n",
      "- prediction_loss_only \n",
      "  default: False \n",
      "  metadata_dict: {'help': 'When performing evaluation and predictions, only returns the loss.'} \n",
      "\n",
      "- per_device_train_batch_size \n",
      "  default: 8 \n",
      "  metadata_dict: {'help': 'Batch size per GPU/TPU core/CPU for training.'} \n",
      "\n",
      "- per_device_eval_batch_size \n",
      "  default: 8 \n",
      "  metadata_dict: {'help': 'Batch size per GPU/TPU core/CPU for evaluation.'} \n",
      "\n",
      "- per_gpu_train_batch_size \n",
      "  default: None \n",
      "  metadata_dict: {'help': 'Deprecated, the use of `--per_device_train_batch_size` is preferred. Batch size per GPU/TPU core/CPU for training.'} \n",
      "\n",
      "- per_gpu_eval_batch_size \n",
      "  default: None \n",
      "  metadata_dict: {'help': 'Deprecated, the use of `--per_device_eval_batch_size` is preferred. Batch size per GPU/TPU core/CPU for evaluation.'} \n",
      "\n",
      "- gradient_accumulation_steps \n",
      "  default: 1 \n",
      "  metadata_dict: {'help': 'Number of updates steps to accumulate before performing a backward/update pass.'} \n",
      "\n",
      "- eval_accumulation_steps \n",
      "  default: None \n",
      "  metadata_dict: {'help': 'Number of predictions steps to accumulate before moving the tensors to the CPU.'} \n",
      "\n",
      "- eval_delay \n",
      "  default: 0 \n",
      "  metadata_dict: {'help': 'Number of epochs or steps to wait for before the first evaluation can be performed, depending on the evaluation_strategy.'} \n",
      "\n",
      "- learning_rate \n",
      "  default: 5e-05 \n",
      "  metadata_dict: {'help': 'The initial learning rate for AdamW.'} \n",
      "\n",
      "- weight_decay \n",
      "  default: 0.0 \n",
      "  metadata_dict: {'help': 'Weight decay for AdamW if we apply some.'} \n",
      "\n",
      "- adam_beta1 \n",
      "  default: 0.9 \n",
      "  metadata_dict: {'help': 'Beta1 for AdamW optimizer'} \n",
      "\n",
      "- adam_beta2 \n",
      "  default: 0.999 \n",
      "  metadata_dict: {'help': 'Beta2 for AdamW optimizer'} \n",
      "\n",
      "- adam_epsilon \n",
      "  default: 1e-08 \n",
      "  metadata_dict: {'help': 'Epsilon for AdamW optimizer.'} \n",
      "\n",
      "- max_grad_norm \n",
      "  default: 1.0 \n",
      "  metadata_dict: {'help': 'Max gradient norm.'} \n",
      "\n",
      "- num_train_epochs \n",
      "  default: 3.0 \n",
      "  metadata_dict: {'help': 'Total number of training epochs to perform.'} \n",
      "\n",
      "- max_steps \n",
      "  default: -1 \n",
      "  metadata_dict: {'help': 'If > 0: set total number of training steps to perform. Override num_train_epochs.'} \n",
      "\n",
      "- lr_scheduler_type \n",
      "  default: linear \n",
      "  metadata_dict: {'help': 'The scheduler type to use.'} \n",
      "\n",
      "- warmup_ratio \n",
      "  default: 0.0 \n",
      "  metadata_dict: {'help': 'Linear warmup over warmup_ratio fraction of total steps.'} \n",
      "\n",
      "- warmup_steps \n",
      "  default: 0 \n",
      "  metadata_dict: {'help': 'Linear warmup over warmup_steps.'} \n",
      "\n",
      "- log_level \n",
      "  default: passive \n",
      "  metadata_dict: {'help': \"Logger log level to use on the main node. Possible choices are the log levels as strings: 'debug', 'info', 'warning', 'error' and 'critical', plus a 'passive' level which doesn't set anything and lets the application set the level. Defaults to 'passive'.\", 'choices': dict_keys(['debug', 'info', 'warning', 'error', 'critical', 'passive'])} \n",
      "\n",
      "- log_level_replica \n",
      "  default: passive \n",
      "  metadata_dict: {'help': 'Logger log level to use on replica nodes. Same choices and defaults as ``log_level``', 'choices': dict_keys(['debug', 'info', 'warning', 'error', 'critical', 'passive'])} \n",
      "\n",
      "- log_on_each_node \n",
      "  default: True \n",
      "  metadata_dict: {'help': 'When doing a multinode distributed training, whether to log once per node or just once on the main node.'} \n",
      "\n",
      "- logging_dir \n",
      "  default: None \n",
      "  metadata_dict: {'help': 'Tensorboard log dir.'} \n",
      "\n",
      "- logging_strategy \n",
      "  default: steps \n",
      "  metadata_dict: {'help': 'The logging strategy to use.'} \n",
      "\n",
      "- logging_first_step \n",
      "  default: False \n",
      "  metadata_dict: {'help': 'Log the first global_step'} \n",
      "\n",
      "- logging_steps \n",
      "  default: 500 \n",
      "  metadata_dict: {'help': 'Log every X updates steps.'} \n",
      "\n",
      "- logging_nan_inf_filter \n",
      "  default: True \n",
      "  metadata_dict: {'help': 'Filter nan and inf losses for logging.'} \n",
      "\n",
      "- save_strategy \n",
      "  default: steps \n",
      "  metadata_dict: {'help': 'The checkpoint save strategy to use.'} \n",
      "\n",
      "- save_steps \n",
      "  default: 500 \n",
      "  metadata_dict: {'help': 'Save checkpoint every X updates steps.'} \n",
      "\n",
      "- save_total_limit \n",
      "  default: None \n",
      "  metadata_dict: {'help': 'Limit the total amount of checkpoints. Deletes the older checkpoints in the output_dir. Default is unlimited checkpoints'} \n",
      "\n",
      "- save_on_each_node \n",
      "  default: False \n",
      "  metadata_dict: {'help': 'When doing multi-node distributed training, whether to save models and checkpoints on each node, or only on the main one'} \n",
      "\n",
      "- no_cuda \n",
      "  default: False \n",
      "  metadata_dict: {'help': 'Do not use CUDA even when it is available'} \n",
      "\n",
      "- use_mps_device \n",
      "  default: False \n",
      "  metadata_dict: {'help': 'Whether to use Apple Silicon chip based `mps` device.'} \n",
      "\n",
      "- seed \n",
      "  default: 42 \n",
      "  metadata_dict: {'help': 'Random seed that will be set at the beginning of training.'} \n",
      "\n",
      "- data_seed \n",
      "  default: None \n",
      "  metadata_dict: {'help': 'Random seed to be used with data samplers.'} \n",
      "\n",
      "- jit_mode_eval \n",
      "  default: False \n",
      "  metadata_dict: {'help': 'Whether or not to use PyTorch jit trace for inference'} \n",
      "\n",
      "- use_ipex \n",
      "  default: False \n",
      "  metadata_dict: {'help': \"Use Intel extension for PyTorch when it is available, installation: 'https://github.com/intel/intel-extension-for-pytorch'\"} \n",
      "\n",
      "- bf16 \n",
      "  default: False \n",
      "  metadata_dict: {'help': 'Whether to use bf16 (mixed) precision instead of 32-bit. Requires Ampere or higher NVIDIA architecture or using CPU (no_cuda). This is an experimental API and it may change.'} \n",
      "\n",
      "- fp16 \n",
      "  default: False \n",
      "  metadata_dict: {'help': 'Whether to use fp16 (mixed) precision instead of 32-bit'} \n",
      "\n",
      "- fp16_opt_level \n",
      "  default: O1 \n",
      "  metadata_dict: {'help': \"For fp16: Apex AMP optimization level selected in ['O0', 'O1', 'O2', and 'O3']. See details at https://nvidia.github.io/apex/amp.html\"} \n",
      "\n",
      "- half_precision_backend \n",
      "  default: auto \n",
      "  metadata_dict: {'help': 'The backend to be used for half precision.', 'choices': ['auto', 'cuda_amp', 'apex', 'cpu_amp']} \n",
      "\n",
      "- bf16_full_eval \n",
      "  default: False \n",
      "  metadata_dict: {'help': 'Whether to use full bfloat16 evaluation instead of 32-bit. This is an experimental API and it may change.'} \n",
      "\n",
      "- fp16_full_eval \n",
      "  default: False \n",
      "  metadata_dict: {'help': 'Whether to use full float16 evaluation instead of 32-bit'} \n",
      "\n",
      "- tf32 \n",
      "  default: None \n",
      "  metadata_dict: {'help': 'Whether to enable tf32 mode, available in Ampere and newer GPU architectures. This is an experimental API and it may change.'} \n",
      "\n",
      "- local_rank \n",
      "  default: -1 \n",
      "  metadata_dict: {'help': 'For distributed training: local_rank'} \n",
      "\n",
      "- xpu_backend \n",
      "  default: None \n",
      "  metadata_dict: {'help': 'The backend to be used for distributed training on Intel XPU.', 'choices': ['mpi', 'ccl', 'gloo']} \n",
      "\n",
      "- tpu_num_cores \n",
      "  default: None \n",
      "  metadata_dict: {'help': 'TPU: Number of TPU cores (automatically passed by launcher script)'} \n",
      "\n",
      "- tpu_metrics_debug \n",
      "  default: False \n",
      "  metadata_dict: {'help': 'Deprecated, the use of `--debug tpu_metrics_debug` is preferred. TPU: Whether to print debug metrics'} \n",
      "\n",
      "- debug \n",
      "  default:  \n",
      "  metadata_dict: {'help': 'Whether or not to enable debug mode. Current options: `underflow_overflow` (Detect underflow and overflow in activations and weights), `tpu_metrics_debug` (print debug metrics on TPU).'} \n",
      "\n",
      "- dataloader_drop_last \n",
      "  default: False \n",
      "  metadata_dict: {'help': 'Drop the last incomplete batch if it is not divisible by the batch size.'} \n",
      "\n",
      "- eval_steps \n",
      "  default: None \n",
      "  metadata_dict: {'help': 'Run an evaluation every X steps.'} \n",
      "\n",
      "- dataloader_num_workers \n",
      "  default: 0 \n",
      "  metadata_dict: {'help': 'Number of subprocesses to use for data loading (PyTorch only). 0 means that the data will be loaded in the main process.'} \n",
      "\n",
      "- past_index \n",
      "  default: -1 \n",
      "  metadata_dict: {'help': 'If >=0, uses the corresponding part of the output as the past state for next step.'} \n",
      "\n",
      "- run_name \n",
      "  default: None \n",
      "  metadata_dict: {'help': 'An optional descriptor for the run. Notably used for wandb logging.'} \n",
      "\n",
      "- disable_tqdm \n",
      "  default: None \n",
      "  metadata_dict: {'help': 'Whether or not to disable the tqdm progress bars.'} \n",
      "\n",
      "- remove_unused_columns \n",
      "  default: True \n",
      "  metadata_dict: {'help': 'Remove columns not required by the model when using an nlp.Dataset.'} \n",
      "\n",
      "- label_names \n",
      "  default: None \n",
      "  metadata_dict: {'help': 'The list of keys in your dictionary of inputs that correspond to the labels.'} \n",
      "\n",
      "- load_best_model_at_end \n",
      "  default: False \n",
      "  metadata_dict: {'help': 'Whether or not to load the best model found during training at the end of training.'} \n",
      "\n",
      "- metric_for_best_model \n",
      "  default: None \n",
      "  metadata_dict: {'help': 'The metric to use to compare two different models.'} \n",
      "\n",
      "- greater_is_better \n",
      "  default: None \n",
      "  metadata_dict: {'help': 'Whether the `metric_for_best_model` should be maximized or not.'} \n",
      "\n",
      "- ignore_data_skip \n",
      "  default: False \n",
      "  metadata_dict: {'help': 'When resuming training, whether or not to skip the first epochs and batches to get to the same training data.'} \n",
      "\n",
      "- sharded_ddp \n",
      "  default:  \n",
      "  metadata_dict: {'help': 'Whether or not to use sharded DDP training (in distributed training only). The base option should be `simple`, `zero_dp_2` or `zero_dp_3` and you can add CPU-offload to `zero_dp_2` or `zero_dp_3` like this: zero_dp_2 offload` or `zero_dp_3 offload`. You can add auto-wrap to `zero_dp_2` or `zero_dp_3` with the same syntax: zero_dp_2 auto_wrap` or `zero_dp_3 auto_wrap`.'} \n",
      "\n",
      "- fsdp \n",
      "  default:  \n",
      "  metadata_dict: {'help': 'Whether or not to use PyTorch Fully Sharded Data Parallel (FSDP) training (in distributed training only). The base option should be `full_shard`, `shard_grad_op` or `no_shard` and you can add CPU-offload to `full_shard` or `shard_grad_op` like this: full_shard offload` or `shard_grad_op offload`. You can add auto-wrap to `full_shard` or `shard_grad_op` with the same syntax: full_shard auto_wrap` or `shard_grad_op auto_wrap`.'} \n",
      "\n",
      "- fsdp_min_num_params \n",
      "  default: 0 \n",
      "  metadata_dict: {'help': \"FSDP's minimum number of parameters for Default Auto Wrapping. (useful only when `fsdp` field is passed).\"} \n",
      "\n",
      "- fsdp_transformer_layer_cls_to_wrap \n",
      "  default: None \n",
      "  metadata_dict: {'help': 'Transformer layer class name (case-sensitive) to wrap ,e.g, `BertLayer`, `GPTJBlock`, `T5Block` .... (useful only when `fsdp` flag is passed).'} \n",
      "\n",
      "- deepspeed \n",
      "  default: None \n",
      "  metadata_dict: {'help': 'Enable deepspeed and pass the path to deepspeed json config file (e.g. ds_config.json) or an already loaded json file as a dict'} \n",
      "\n",
      "- label_smoothing_factor \n",
      "  default: 0.0 \n",
      "  metadata_dict: {'help': 'The label smoothing epsilon to apply (zero means no label smoothing).'} \n",
      "\n",
      "- optim \n",
      "  default: adamw_hf \n",
      "  metadata_dict: {'help': 'The optimizer to use.'} \n",
      "\n",
      "- optim_args \n",
      "  default: None \n",
      "  metadata_dict: {'help': 'Optional arguments to supply to optimizer.'} \n",
      "\n",
      "- adafactor \n",
      "  default: False \n",
      "  metadata_dict: {'help': 'Whether or not to replace AdamW by Adafactor.'} \n",
      "\n",
      "- group_by_length \n",
      "  default: False \n",
      "  metadata_dict: {'help': 'Whether or not to group samples of roughly the same length together when batching.'} \n",
      "\n",
      "- length_column_name \n",
      "  default: length \n",
      "  metadata_dict: {'help': 'Column name with precomputed lengths to use when grouping by length.'} \n",
      "\n",
      "- report_to \n",
      "  default: None \n",
      "  metadata_dict: {'help': 'The list of integrations to report the results and logs to.'} \n",
      "\n",
      "- ddp_find_unused_parameters \n",
      "  default: None \n",
      "  metadata_dict: {'help': 'When using distributed training, the value of the flag `find_unused_parameters` passed to `DistributedDataParallel`.'} \n",
      "\n",
      "- ddp_bucket_cap_mb \n",
      "  default: None \n",
      "  metadata_dict: {'help': 'When using distributed training, the value of the flag `bucket_cap_mb` passed to `DistributedDataParallel`.'} \n",
      "\n",
      "- dataloader_pin_memory \n",
      "  default: True \n",
      "  metadata_dict: {'help': 'Whether or not to pin memory for DataLoader.'} \n",
      "\n",
      "- skip_memory_metrics \n",
      "  default: True \n",
      "  metadata_dict: {'help': 'Whether or not to skip adding of memory profiler reports to metrics.'} \n",
      "\n",
      "- use_legacy_prediction_loop \n",
      "  default: False \n",
      "  metadata_dict: {'help': 'Whether or not to use the legacy prediction_loop in the Trainer.'} \n",
      "\n",
      "- push_to_hub \n",
      "  default: False \n",
      "  metadata_dict: {'help': 'Whether or not to upload the trained model to the model hub after training.'} \n",
      "\n",
      "- resume_from_checkpoint \n",
      "  default: None \n",
      "  metadata_dict: {'help': 'The path to a folder with a valid checkpoint for your model.'} \n",
      "\n",
      "- hub_model_id \n",
      "  default: None \n",
      "  metadata_dict: {'help': 'The name of the repository to keep in sync with the local `output_dir`.'} \n",
      "\n",
      "- hub_strategy \n",
      "  default: every_save \n",
      "  metadata_dict: {'help': 'The hub strategy to use when `--push_to_hub` is activated.'} \n",
      "\n",
      "- hub_token \n",
      "  default: None \n",
      "  metadata_dict: {'help': 'The token to use to push to the Model Hub.'} \n",
      "\n",
      "- hub_private_repo \n",
      "  default: False \n",
      "  metadata_dict: {'help': 'Whether the model repository is private or not.'} \n",
      "\n",
      "- gradient_checkpointing \n",
      "  default: False \n",
      "  metadata_dict: {'help': 'If True, use gradient checkpointing to save memory at the expense of slower backward pass.'} \n",
      "\n",
      "- include_inputs_for_metrics \n",
      "  default: False \n",
      "  metadata_dict: {'help': 'Whether or not the inputs will be passed to the `compute_metrics` function.'} \n",
      "\n",
      "- fp16_backend \n",
      "  default: auto \n",
      "  metadata_dict: {'help': 'Deprecated. Use half_precision_backend instead', 'choices': ['auto', 'cuda_amp', 'apex', 'cpu_amp']} \n",
      "\n",
      "- push_to_hub_model_id \n",
      "  default: None \n",
      "  metadata_dict: {'help': 'The name of the repository to which push the `Trainer`.'} \n",
      "\n",
      "- push_to_hub_organization \n",
      "  default: None \n",
      "  metadata_dict: {'help': 'The name of the organization in with to which push the `Trainer`.'} \n",
      "\n",
      "- push_to_hub_token \n",
      "  default: None \n",
      "  metadata_dict: {'help': 'The token to use to push to the Model Hub.'} \n",
      "\n",
      "- _n_gpu \n",
      "  default: -1 \n",
      "  metadata_dict: {} \n",
      "\n",
      "- mp_parameters \n",
      "  default:  \n",
      "  metadata_dict: {'help': 'Used by the SageMaker launcher to send mp-specific args. Ignored in Trainer'} \n",
      "\n",
      "- auto_find_batch_size \n",
      "  default: False \n",
      "  metadata_dict: {'help': 'Whether to automatically decrease the batch size in half and rerun the training loop again each time a CUDA Out-of-Memory was reached'} \n",
      "\n",
      "- full_determinism \n",
      "  default: False \n",
      "  metadata_dict: {'help': 'Whether to call enable_full_determinism instead of set_seed for reproducibility in distributed training'} \n",
      "\n",
      "- torchdynamo \n",
      "  default: None \n",
      "  metadata_dict: {'help': 'Sets up the backend compiler for TorchDynamo.', 'choices': ['eager', 'aot_eager', 'inductor', 'nvfuser', 'aot_nvfuser', 'aot_cudagraphs', 'ofi', 'fx2trt', 'onnxrt', 'ipex']} \n",
      "\n",
      "- ray_scope \n",
      "  default: last \n",
      "  metadata_dict: {'help': 'The scope to use when doing hyperparameter search with Ray. By default, `\"last\"` will be used. Ray will then use the last checkpoint of all trials, compare those, and select the best one. However, other options are also available. See the Ray documentation (https://docs.ray.io/en/latest/tune/api_docs/analysis.html#ray.tune.ExperimentAnalysis.get_best_trial) for more options.'} \n",
      "\n",
      "- ddp_timeout \n",
      "  default: 1800 \n",
      "  metadata_dict: {'help': 'Overrides the default timeout for distributed training (value should be given in seconds).'} \n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2022-12-20 09:05:32,902 - 140132563211072 - info.py-info:365 - INFO: Loading Dataset Infos from /home/zhangchen/.cache/huggingface/modules/datasets_modules/datasets/wikitext/a241db52902eaf2c6aa732210bead40c090019a499ceb13bcbfa3f8ab646a126\n",
      "2022-12-20 09:05:32,916 - 140132563211072 - builder.py-builder:354 - INFO: Overwrite dataset info from restored data version.\n",
      "2022-12-20 09:05:32,917 - 140132563211072 - info.py-info:285 - INFO: Loading Dataset info from /home/zhangchen/.cache/huggingface/datasets/wikitext/wikitext-2-raw-v1/1.0.0/a241db52902eaf2c6aa732210bead40c090019a499ceb13bcbfa3f8ab646a126\n",
      "2022-12-20 09:05:32,932 - 140132563211072 - builder.py-builder:747 - WARNING: Found cached dataset wikitext (/home/zhangchen/.cache/huggingface/datasets/wikitext/wikitext-2-raw-v1/1.0.0/a241db52902eaf2c6aa732210bead40c090019a499ceb13bcbfa3f8ab646a126)\n",
      "2022-12-20 09:05:32,933 - 140132563211072 - info.py-info:285 - INFO: Loading Dataset info from /home/zhangchen/.cache/huggingface/datasets/wikitext/wikitext-2-raw-v1/1.0.0/a241db52902eaf2c6aa732210bead40c090019a499ceb13bcbfa3f8ab646a126\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "525fd4174bda43d391f4d7c022de2faf",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/3 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2022-12-20 09:05:32,966 - 140132563211072 - arrow_dataset.py-arrow_dataset:2806 - WARNING: Loading cached processed dataset at /home/zhangchen/.cache/huggingface/datasets/wikitext/wikitext-2-raw-v1/1.0.0/a241db52902eaf2c6aa732210bead40c090019a499ceb13bcbfa3f8ab646a126/cache-fe7ac049e31bb4d9.arrow\n",
      "2022-12-20 09:05:32,975 - 140132563211072 - arrow_dataset.py-arrow_dataset:2806 - WARNING: Loading cached processed dataset at /home/zhangchen/.cache/huggingface/datasets/wikitext/wikitext-2-raw-v1/1.0.0/a241db52902eaf2c6aa732210bead40c090019a499ceb13bcbfa3f8ab646a126/cache-b1f05a932b1ce06d.arrow\n",
      "2022-12-20 09:05:32,984 - 140132563211072 - arrow_dataset.py-arrow_dataset:2806 - WARNING: Loading cached processed dataset at /home/zhangchen/.cache/huggingface/datasets/wikitext/wikitext-2-raw-v1/1.0.0/a241db52902eaf2c6aa732210bead40c090019a499ceb13bcbfa3f8ab646a126/cache-62104dcf4bc00fe6.arrow\n",
      "2022-12-20 09:05:32,991 - 140132563211072 - arrow_dataset.py-arrow_dataset:2806 - WARNING: Loading cached processed dataset at /home/zhangchen/.cache/huggingface/datasets/wikitext/wikitext-2-raw-v1/1.0.0/a241db52902eaf2c6aa732210bead40c090019a499ceb13bcbfa3f8ab646a126/cache-976c2d1bd8ec1d8a.arrow\n",
      "2022-12-20 09:05:32,999 - 140132563211072 - arrow_dataset.py-arrow_dataset:2806 - WARNING: Loading cached processed dataset at /home/zhangchen/.cache/huggingface/datasets/wikitext/wikitext-2-raw-v1/1.0.0/a241db52902eaf2c6aa732210bead40c090019a499ceb13bcbfa3f8ab646a126/cache-6fbb62a735bd0cb8.arrow\n",
      "2022-12-20 09:05:33,007 - 140132563211072 - arrow_dataset.py-arrow_dataset:2806 - WARNING: Loading cached processed dataset at /home/zhangchen/.cache/huggingface/datasets/wikitext/wikitext-2-raw-v1/1.0.0/a241db52902eaf2c6aa732210bead40c090019a499ceb13bcbfa3f8ab646a126/cache-8a07e1c6f15c4f9a.arrow\n",
      "[INFO|trainer.py:703] 2022-12-20 09:05:35,787 >> The following columns in the training set don't have a corresponding argument in `BertForMaskedLM.forward` and have been ignored: special_tokens_mask. If special_tokens_mask are not expected by `BertForMaskedLM.forward`,  you can safely ignore this message.\n",
      "[INFO|trainer.py:1634] 2022-12-20 09:05:35,798 >> ***** Running training *****\n",
      "[INFO|trainer.py:1635] 2022-12-20 09:05:35,799 >>   Num examples = 4627\n",
      "[INFO|trainer.py:1636] 2022-12-20 09:05:35,800 >>   Num Epochs = 3\n",
      "[INFO|trainer.py:1637] 2022-12-20 09:05:35,800 >>   Instantaneous batch size per device = 8\n",
      "[INFO|trainer.py:1638] 2022-12-20 09:05:35,801 >>   Total train batch size (w. parallel, distributed & accumulation) = 64\n",
      "[INFO|trainer.py:1639] 2022-12-20 09:05:35,802 >>   Gradient Accumulation steps = 1\n",
      "[INFO|trainer.py:1640] 2022-12-20 09:05:35,802 >>   Total optimization steps = 219\n",
      "[INFO|trainer.py:1641] 2022-12-20 09:05:35,803 >>   Number of trainable parameters = 109514298\n",
      "[WARNING|logging.py:281] 2022-12-20 09:05:35,844 >> You're using a BertTokenizerFast tokenizer. Please note that with a fast tokenizer, using the `__call__` method is faster than using a method to encode the text followed by a call to the `pad` method to get a padded encoding.\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "\n",
       "    <div>\n",
       "      \n",
       "      <progress value='219' max='219' style='width:300px; height:20px; vertical-align: middle;'></progress>\n",
       "      [219/219 01:15, Epoch 3/3]\n",
       "    </div>\n",
       "    <table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       " <tr style=\"text-align: left;\">\n",
       "      <th>Step</th>\n",
       "      <th>Training Loss</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "  </tbody>\n",
       "</table><p>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[INFO|trainer.py:1885] 2022-12-20 09:07:25,052 >> \n",
      "\n",
      "Training completed. Do not forget to share your model on huggingface.co/models =)\n",
      "\n",
      "\n",
      "[INFO|trainer.py:2693] 2022-12-20 09:07:25,055 >> Saving model checkpoint to ./tmp/test-mlm\n",
      "[INFO|configuration_utils.py:447] 2022-12-20 09:07:25,057 >> Configuration saved in ./tmp/test-mlm/config.json\n",
      "[INFO|modeling_utils.py:1637] 2022-12-20 09:07:25,790 >> Model weights saved in ./tmp/test-mlm/pytorch_model.bin\n",
      "[INFO|tokenization_utils_base.py:2157] 2022-12-20 09:07:25,792 >> tokenizer config file saved in ./tmp/test-mlm/tokenizer_config.json\n",
      "[INFO|tokenization_utils_base.py:2164] 2022-12-20 09:07:25,794 >> Special tokens file saved in ./tmp/test-mlm/special_tokens_map.json\n",
      "2022-12-20 09:07:25,833 - 140132563211072 - train_mlm_with_hf_trainer.py-train_mlm_with_hf_trainer:443 - INFO: *** Evaluate ***\n",
      "[INFO|trainer.py:703] 2022-12-20 09:07:25,834 >> The following columns in the evaluation set don't have a corresponding argument in `BertForMaskedLM.forward` and have been ignored: special_tokens_mask. If special_tokens_mask are not expected by `BertForMaskedLM.forward`,  you can safely ignore this message.\n",
      "[INFO|trainer.py:2944] 2022-12-20 09:07:25,837 >> ***** Running Evaluation *****\n",
      "[INFO|trainer.py:2946] 2022-12-20 09:07:25,837 >>   Num examples = 479\n",
      "[INFO|trainer.py:2949] 2022-12-20 09:07:25,838 >>   Batch size = 64\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "***** train metrics *****\n",
      "  epoch                    =        3.0\n",
      "  total_flos               =  3402629GF\n",
      "  train_loss               =     1.8067\n",
      "  train_runtime            = 0:01:49.24\n",
      "  train_samples            =       4627\n",
      "  train_samples_per_second =    127.058\n",
      "  train_steps_per_second   =      2.005\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "\n",
       "    <div>\n",
       "      \n",
       "      <progress value='8' max='8' style='width:300px; height:20px; vertical-align: middle;'></progress>\n",
       "      [8/8 00:01]\n",
       "    </div>\n",
       "    "
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "***** eval metrics *****\n",
      "  epoch                   =        3.0\n",
      "  eval_accuracy           =     0.6789\n",
      "  eval_loss               =      1.575\n",
      "  eval_runtime            = 0:00:01.89\n",
      "  eval_samples            =        479\n",
      "  eval_samples_per_second =    252.993\n",
      "  eval_steps_per_second   =      4.225\n",
      "  perplexity              =     4.8309\n",
      "done mlm.\n"
     ]
    }
   ],
   "source": [
    "bert_op.train(task='mlm', data_args=data_args, training_args=training_args)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "be8c97e6",
   "metadata": {},
   "source": [
    "`data_args` specifies the training data set, you can specify the name directly, it will download the specified dataset through [datasets](https://huggingface.co/docs/datasets/index). For more `data_args` infos, refer to the `**** DataTrainingArguments ****` line in this block's output.  \n",
    "\n",
    "`training_args` specifies the training config using [transformer TrainingArguments](https://huggingface.co/docs/transformers/main/en/main_classes/trainer#transformers.TrainingArguments). For more `training_args` infos, refer to the `**** TrainingArguments ****` line in this block's output."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "db793daf",
   "metadata": {},
   "source": [
    "If you run this script for the first time, it will automatically download the corresponding model and data set, which may require at least 1.5G space."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4f67c83f",
   "metadata": {},
   "source": [
    "If you run this training process very slowly, please make sure your device contains advanced GPUs, or, you may reduce the number of epochs in order to simply run through the training. By default, training will be performed in parallel using available GPUs."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cce1a9fc",
   "metadata": {},
   "source": [
    "If you see the final evaluate result with `eval_accuracy` about 0.65, it means you have successfully trained the operator. And the output result with model's weights has been in your `output_dir`."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3356172f",
   "metadata": {},
   "source": [
    "### Use your fine-tuned weights\n",
    "Please note that the model trained in this way is a model with a specific task header, but the original operator does not contain the head. If you need to use your trained weights to extract embedding, you need to convert model weights to adapt model without head and load it. Here is a example."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "dc7831ff",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from collections import OrderedDict\n",
    "from transformers import BertModel\n",
    "from transformers.utils import logging\n",
    "\n",
    "logging.set_verbosity_error()\n",
    "\n",
    "def convert_bert_weights(trained_weights_path, new_weight_path):\n",
    "    state_dict = torch.load(trained_weights_path, map_location='cpu')\n",
    "    new_state_dict = OrderedDict()\n",
    "    for k, v in state_dict.items():\n",
    "        if k.startswith('bert.'):\n",
    "            new_k = k[5:]\n",
    "            new_state_dict[new_k] = v\n",
    "    bert_model = BertModel.from_pretrained(\"bert-base-uncased\")\n",
    "    bert_model.load_state_dict(new_state_dict, strict=False)\n",
    "    torch.save(bert_model.state_dict(), new_weight_path)\n",
    "\n",
    "convert_bert_weights('./tmp/test-mlm/pytorch_model.bin', './tmp/test-mlm/bert_without_head_weight.bin')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "c5672f1e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...\n",
      "To disable this warning, you can either:\n",
      "\t- Avoid using `tokenizers` before the fork if possible\n",
      "\t- Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)\n"
     ]
    }
   ],
   "source": [
    "new_bert_op = towhee.ops.text_embedding.transformers(\n",
    "    model_name='bert-base-uncased', \n",
    "    checkpoint_path='./tmp/test-mlm/bert_without_head_weight.bin'\n",
    ").get_op()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "83f3b3ef",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(array([[-0.5528104 ,  0.15354109,  0.27865022, ..., -0.42180687,\n",
       "         -0.17302679,  0.84261256],\n",
       "        [-1.0168282 ,  0.6152043 ,  1.5796373 , ..., -0.066678  ,\n",
       "          0.01646271,  1.0581956 ],\n",
       "        [ 0.47564444,  0.3604667 ,  0.87738824, ..., -0.72136706,\n",
       "          0.0545284 ,  1.0028253 ],\n",
       "        ...,\n",
       "        [-0.6695933 ,  0.1092175 ,  1.1711102 , ...,  0.12439835,\n",
       "         -0.6637869 ,  1.0748045 ],\n",
       "        [-0.7066035 , -0.12123797,  0.14397985, ...,  0.7948587 ,\n",
       "          0.06388341,  0.16218108],\n",
       "        [ 0.88073194,  0.20878866, -0.11732787, ...,  0.22331028,\n",
       "         -0.5018618 , -0.35616544]], dtype=float32),\n",
       " (9, 768))"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "embedding = new_bert_op('Hello world, hello every one.')\n",
    "embedding, embedding.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "eeff6a82",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}